#include "LumineRenderer.hpp"
#include "Core/Context/LumineContext.hpp"
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <chrono>
#include "Core/Buffer/BufferManager.hpp"
#include "Core/Texture/TextureManager.hpp"

void Lumine::LumineRenderer::createCommandPool()
{
    vk::CommandPoolCreateInfo info;
    info.setFlags(vk::CommandPoolCreateFlagBits::eResetCommandBuffer)
        .setQueueFamilyIndex(LumineContext::Get().GraphicIndex());
    m_commandPool = vk::raii::CommandPool(LumineContext::Get().Device(), info);
}

void Lumine::LumineRenderer::createCommandBuffers()
{
    m_commandBuffers.clear();
    vk::CommandBufferAllocateInfo allocInfo;
    allocInfo.setCommandPool(m_commandPool)
        .setLevel(vk::CommandBufferLevel::ePrimary)
        .setCommandBufferCount(MAX_FRAMES_IN_FLIGHT);
    m_commandBuffers = std::move(vk::raii::CommandBuffers(LumineContext::Get().Device(), allocInfo));
}

void Lumine::LumineRenderer::createSyncObjects()
{
    imageAvailableSemaphores.clear();
    renderFinishedSemaphores.clear();
    inFlightFences.clear();
    vk::FenceCreateInfo info;
    info.setFlags(vk::FenceCreateFlagBits::eSignaled);
    for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        inFlightFences.emplace_back(LumineContext::Get().Device(), info);
    }
    for (size_t i = 0; i < m_swapchainRef->SwapchainImageViews().size(); i++)
    {
        imageAvailableSemaphores.emplace_back(LumineContext::Get().Device(), vk::SemaphoreCreateInfo());
        renderFinishedSemaphores.emplace_back(LumineContext::Get().Device(), vk::SemaphoreCreateInfo());
    }
}

void Lumine::LumineRenderer::createDescriptorPool()
{
    std::array poolSize{
        vk::DescriptorPoolSize(vk::DescriptorType::eUniformBuffer, MAX_FRAMES_IN_FLIGHT),
        vk::DescriptorPoolSize(vk::DescriptorType::eCombinedImageSampler, MAX_FRAMES_IN_FLIGHT),
    };
    vk::DescriptorPoolCreateInfo poolInfo;
    poolInfo.setFlags(vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet)
        .setMaxSets(MAX_FRAMES_IN_FLIGHT)
        .setPoolSizeCount(poolSize.size())
        .setPPoolSizes(poolSize.data());
    descriptorPool = vk::raii::DescriptorPool(LumineContext::Get().Device(), poolInfo);
}

void Lumine::LumineRenderer::createDescriptorSets()
{
    std::vector<vk::DescriptorSetLayout> layouts(MAX_FRAMES_IN_FLIGHT, m_pipelineRef->DescriptorSetLayout());
    vk::DescriptorSetAllocateInfo allocInfo;
    allocInfo.setDescriptorPool(descriptorPool)
        .setDescriptorSetCount(static_cast<uint32_t>(layouts.size()))
        .setPSetLayouts(layouts.data());
    descriptorSets.clear();
    descriptorSets = LumineContext::Get().Device().allocateDescriptorSets(allocInfo);
    for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
    {
        vk::DescriptorBufferInfo bufferInfo;
        bufferInfo.setBuffer(BufferManager::uniformBuffers.GetUniformBuffer(i))
            .setOffset(0)
            .setRange(sizeof(UniformBufferObject));
        vk::DescriptorImageInfo imageInfo;
        imageInfo.setSampler(TextureManager::textures[0].Sampler())
            .setImageView(TextureManager::textures[0].ImageView())
            .setImageLayout(vk::ImageLayout::eShaderReadOnlyOptimal);

        vk::WriteDescriptorSet uniformWrite;
        uniformWrite.setDstSet(descriptorSets[i])
            .setDstBinding(0)
            .setDstArrayElement(0)
            .setDescriptorCount(1)
            .setDescriptorType(vk::DescriptorType::eUniformBuffer)
            .setPBufferInfo(&bufferInfo);
        vk::WriteDescriptorSet imageWrite;
        imageWrite.setDstSet(descriptorSets[i])
            .setDstBinding(1)
            .setDstArrayElement(0)
            .setDescriptorCount(1)
            .setDescriptorType(vk::DescriptorType::eCombinedImageSampler)
            .setPImageInfo(&imageInfo);

        std::array descWrites{
            uniformWrite,
            imageWrite};
        LumineContext::Get()
            .Device()
            .updateDescriptorSets(descWrites, {});
    }
}

vk::raii::CommandBuffer Lumine::LumineRenderer::beginSingleTimeCommands()
{
    vk::CommandBufferAllocateInfo allocInfo;
    allocInfo.setCommandPool(m_commandPool)
        .setLevel(vk::CommandBufferLevel::ePrimary)
        .setCommandBufferCount(1);
    vk::raii::CommandBuffer cmdBuffer = std::move(LumineContext::Get().Device().allocateCommandBuffers(allocInfo).front());

    vk::CommandBufferBeginInfo beginInfo;
    beginInfo.setFlags(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);
    cmdBuffer.begin(beginInfo);
    return cmdBuffer;
}

void Lumine::LumineRenderer::endSingleTimeCommands(vk::raii::CommandBuffer &commandBuffer)
{
    commandBuffer.end();
    vk::SubmitInfo submitInfo;
    submitInfo.setCommandBufferCount(1)
        .setPCommandBuffers(&*commandBuffer);
    LumineContext::Get().GraphicQueue().submit(submitInfo);
    LumineContext::Get().GraphicQueue().waitIdle();
}

void Lumine::LumineRenderer::recordCommandBuffer(vk::raii::CommandBuffer &commandBuffer, uint32_t imageIndex)
{

    commandBuffer.begin({});
    vk::ClearValue clearColor = vk::ClearColorValue(0.0f, 0.0f, 0.0f, 1.0f);
    transition_image_layout(commandBuffer,
                            imageIndex,
                            vk::ImageLayout::eUndefined,
                            vk::ImageLayout::eColorAttachmentOptimal,
                            {},
                            vk::AccessFlagBits2::eColorAttachmentWrite,
                            vk::PipelineStageFlagBits2::eTopOfPipe,
                            vk::PipelineStageFlagBits2::eColorAttachmentOutput);
    // 设置图像附件
    vk::RenderingAttachmentInfo info;
    info.setImageView(*m_swapchainRef->SwapchainImageViews()[imageIndex])
        .setImageLayout(vk::ImageLayout::eColorAttachmentOptimal)
        .setLoadOp(vk::AttachmentLoadOp::eClear)
        .setStoreOp(vk::AttachmentStoreOp::eStore)
        .setClearValue(clearColor);
    vk::RenderingAttachmentInfo depthAttachment{};
    depthAttachment.setImageView(TextureManager::depth[0].ImageView())
        .setImageLayout(vk::ImageLayout::eDepthStencilAttachmentOptimal)
        .setLoadOp(vk::AttachmentLoadOp::eClear)
        .setStoreOp(vk::AttachmentStoreOp::eDontCare)
        .setClearValue(vk::ClearValue(vk::ClearDepthStencilValue(1.0f, 0)));

    vk::RenderingInfo renderInfo;
    vk::RenderingAreaInfoKHR areaInfo;
    renderInfo.setRenderArea({{0, 0}, m_swapchainRef->Extent()})
        .setLayerCount(1)
        .setColorAttachmentCount(1)
        .setPColorAttachments(&info)
        .setPDepthAttachment(&depthAttachment);
    commandBuffer.beginRendering(renderInfo);
    commandBuffer.bindPipeline(vk::PipelineBindPoint::eGraphics, m_pipeline);
    commandBuffer.bindVertexBuffers(0, *BufferManager::vertexBuffers.Buffer(), {0});
    commandBuffer.bindIndexBuffer(*BufferManager::indexBuffers.Buffer(), 0, vk::IndexType::eUint16);
    commandBuffer.setViewport(0, vk::Viewport(0.0f, 0.0f, static_cast<float>(m_swapchainRef->Extent().width), static_cast<float>(m_swapchainRef->Extent().height), 0.0f, 1.0f));
    commandBuffer.setScissor(0, vk::Rect2D(vk::Offset2D(0, 0), m_swapchainRef->Extent()));
    commandBuffer.bindDescriptorSets(vk::PipelineBindPoint::eGraphics, m_pipelineRef->PipelineLayout(), 0, *descriptorSets[currentFrame], nullptr);
    commandBuffer.drawIndexed(squareIndices.size(), 1, 0, 0, 0);
    commandBuffer.endRendering();
    transition_image_layout(commandBuffer, imageIndex,
                            vk::ImageLayout::eColorAttachmentOptimal,
                            vk::ImageLayout::ePresentSrcKHR,
                            vk::AccessFlagBits2::eColorAttachmentWrite,
                            {},
                            vk::PipelineStageFlagBits2::eColorAttachmentOutput,
                            vk::PipelineStageFlagBits2::eBottomOfPipe);
    commandBuffer.end();
}

void Lumine::LumineRenderer::transition_image_layout(vk::raii::CommandBuffer &cmdBuffer, uint32_t imageIndex, vk::ImageLayout oldLayout, vk::ImageLayout newLayout, vk::AccessFlags2 srcAccessMask, vk::AccessFlags2 dstAccessMask, vk::PipelineStageFlags2 srcStageMask, vk::PipelineStageFlags2 dstStageMask)
{
    // --- 优化开始：队列家族所有权转移逻辑 ---
    uint32_t srcQ = VK_QUEUE_FAMILY_IGNORED;
    uint32_t dstQ = VK_QUEUE_FAMILY_IGNORED;
    uint32_t graphicIndex = LumineContext::Get().GraphicIndex();
    uint32_t presentIndex = LumineContext::Get().PresentIndex();
    // 仅在 Graphic 和 Present 队列不同时才需要进行所有权转移
    if (graphicIndex != presentIndex)
    {
        if (oldLayout == vk::ImageLayout::eUndefined && newLayout == vk::ImageLayout::eColorAttachmentOptimal)
        {
            // 第一次转换：从 Present/Acquire 队列 (隐含) 转移到 Graphics 队列
            srcQ = presentIndex;
            dstQ = graphicIndex;
        }
        else if (oldLayout == vk::ImageLayout::eColorAttachmentOptimal && newLayout == vk::ImageLayout::ePresentSrcKHR)
        {
            // 第二次转换：从 Graphics 队列转移到 Present 队列
            srcQ = graphicIndex;
            dstQ = presentIndex;
        }
    }

    vk::ImageMemoryBarrier2 barrier;
    vk::ImageSubresourceRange range;
    range.setAspectMask(vk::ImageAspectFlagBits::eColor)
        .setBaseMipLevel(0)
        .setLevelCount(1)
        .setBaseArrayLayer(0)
        .setLayerCount(1);

    barrier.setSrcStageMask(srcStageMask)
        .setSrcAccessMask(srcAccessMask)
        .setDstStageMask(dstStageMask)
        .setDstAccessMask(dstAccessMask)
        .setOldLayout(oldLayout)
        .setNewLayout(newLayout)
        .setSrcQueueFamilyIndex(srcQ) // ✅ 使用计算出的 srcQ
        .setDstQueueFamilyIndex(dstQ) // ✅ 使用计算出的 dstQ
        .setImage(m_swapchainRef->SwapchainImages()[imageIndex])
        .setSubresourceRange(range);

    vk::DependencyInfo deInfo;
    deInfo.setDependencyFlags({})
        .setImageMemoryBarrierCount(1)
        .setPImageMemoryBarriers(&barrier);

    cmdBuffer.pipelineBarrier2(deInfo);
}

void Lumine::LumineRenderer::Initialize(LumineSwapchain &swapchain, LuminePipeline &pipeline)
{
    m_swapchainRef = &swapchain;
    m_pipelineRef = &pipeline;
    createCommandPool();
    createCommandBuffers();
    createDescriptorPool();
    createDescriptorSets();
    createSyncObjects();
}

vk::Result Lumine::LumineRenderer::drawFrame()
{
    updateUniformBuffer(currentFrame);
    // 1. 获取当前帧的同步对象和命令缓冲区
    vk::raii::Semaphore &imageAvailableSemaphore = imageAvailableSemaphores[currentImage];
    vk::raii::Semaphore &renderFinishedSemaphore = renderFinishedSemaphores[currentImage];
    vk::raii::Fence &inFlightFence = inFlightFences[currentFrame];
    vk::raii::CommandBuffer &commandBuffer = m_commandBuffers[currentFrame];

    // 2. 严格序列化：等待当前帧的围栏，确保上一帧的工作完成
    while (vk::Result::eTimeout == LumineContext::Get().Device().waitForFences(*inFlightFence, vk::True, UINT64_MAX))
    {
    }

    // 3. 重置当前帧的围栏，准备 GPU 提交
    LumineContext::Get().Device().resetFences(*inFlightFence);
    uint32_t imageIndex;
    vk::Result acquireRes;
    // 4. 获取下一个图像（GPU 信号 imageAvailableSemaphore）
    try
    {
        auto [result, index] = m_swapchainRef->Swapchain().acquireNextImage(UINT64_MAX, *imageAvailableSemaphore, nullptr);
        // 5. 处理 Swapchain 重建 (Acquire 失败)
        imageIndex = index;
        acquireRes = result;
        if (result == vk::Result::eErrorOutOfDateKHR)
        {
            currentImage = 0;
            return vk::Result::eErrorOutOfDateKHR; // 立即重建并返回，不执行后续提交
        }
    }
    catch (const vk::OutOfDateKHRError)
    {
        return vk::Result::eErrorOutOfDateKHR;
    }
    // 6. 记录命令
    commandBuffer.reset();
    recordCommandBuffer(commandBuffer, imageIndex);

    // 7. 提交命令到图形队列 (使用 Vulkan 1.3/KHR_synchronization2 结构)

    // a. 等待信号量信息：等待 imageAvailableSemaphore，在颜色附件输出阶段开始前同步
    vk::SemaphoreSubmitInfo waitSemaphoreInfo;
    waitSemaphoreInfo.setSemaphore(*imageAvailableSemaphore)
        .setStageMask(vk::PipelineStageFlagBits2::eColorAttachmentOutput);

    // b. 命令缓冲区信息
    vk::CommandBufferSubmitInfo commandBufferInfo;
    commandBufferInfo.setCommandBuffer(*commandBuffer);

    // c. 信号信号量信息：命令完成后信号 renderFinishedSemaphore
    vk::SemaphoreSubmitInfo signalSemaphoreInfo;
    signalSemaphoreInfo.setSemaphore(*renderFinishedSemaphore)
        .setStageMask(vk::PipelineStageFlagBits2::eAllCommands); // 确保在所有命令完成后信号

    // d. 提交信息 (vk::SubmitInfo2)
    vk::SubmitInfo2 submitInfo2;
    submitInfo2.setWaitSemaphoreInfoCount(1)
        .setPWaitSemaphoreInfos(&waitSemaphoreInfo)
        .setCommandBufferInfoCount(1)
        .setPCommandBufferInfos(&commandBufferInfo)
        .setSignalSemaphoreInfoCount(1)
        .setPSignalSemaphoreInfos(&signalSemaphoreInfo);

    // e. 提交：使用 submit2 提交，并在命令执行完毕后信号 inFlightFence
    LumineContext::Get().GraphicQueue().submit2(submitInfo2, *inFlightFence);

    // 8. 呈现图像
    vk::PresentInfoKHR presentInfo;

    // 等待 renderFinishedSemaphore
    vk::Semaphore signalSemaphores[] = {*renderFinishedSemaphore};
    presentInfo.setWaitSemaphoreCount(1);
    presentInfo.setPWaitSemaphores(signalSemaphores);
    vk::SwapchainKHR scs[] = {*m_swapchainRef->Swapchain()};
    presentInfo.setPSwapchains(scs);
    presentInfo.setImageIndices(imageIndex);
    presentInfo.pResults = nullptr;

    try
    {
        vk::Result result = LumineContext::Get().PresentQueue().presentKHR(presentInfo);
        // 9. 处理 Swapchain 重建（呈现失败/次优）
        // 立即处理这些错误，确保信号量被消耗，从而解决双重信号问题。
        if (result == vk::Result::eErrorOutOfDateKHR || result == vk::Result::eSuboptimalKHR)
        {
            currentImage = 0;
            return result;
        }
        else if (result != vk::Result::eSuccess)
        {
            return result;
        }
    }
    catch (const vk::OutOfDateKHRError)
    {
        return vk::Result::eErrorOutOfDateKHR;
    }

    // 10. 更新帧索引
    currentImage = (currentImage + 1) % m_swapchainRef->SwapchainImageViews().size();
    currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
}

void Lumine::LumineRenderer::copyBuffer(vk::raii::Buffer &srcBuffer, vk::raii::Buffer &dstBuffer, vk::DeviceSize size)
{
    vk::CommandBufferAllocateInfo allocInfo;
    allocInfo.setCommandPool(m_commandPool)
        .setLevel(vk::CommandBufferLevel::ePrimary)
        .setCommandBufferCount(1);
    vk::raii::CommandBuffer cmdCopyBuffer = std::move(LumineContext::Get().Device().allocateCommandBuffers(allocInfo).front());

    vk::CommandBufferBeginInfo cmdBegin;
    cmdBegin.setFlags(vk::CommandBufferUsageFlagBits::eOneTimeSubmit);
    cmdCopyBuffer.begin(cmdBegin);
    cmdCopyBuffer.copyBuffer(srcBuffer, dstBuffer, vk::BufferCopy(0, 0, size));
    cmdCopyBuffer.end();
}

void Lumine::LumineRenderer::updateUniformBuffer(uint32_t currentImage)
{
    static auto startTime = std::chrono::high_resolution_clock::now();
    auto currentTime = std::chrono::high_resolution_clock::now();
    float time = std::chrono::duration<float, std::chrono::seconds::period>(currentTime - startTime).count();
    UniformBufferObject ubo{};
    ubo.model = rotate(glm::mat4(1.0f), time * glm::radians(90.0f), glm::vec3(0.0f, 0.0f, 1.0f));
    ubo.view = lookAt(glm::vec3(2.0f, 2.0f, 2.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 1.0f));
    ubo.proj = glm::perspective(glm::radians(45.0f), static_cast<float>(m_swapchainRef->Extent().width) / static_cast<float>(m_swapchainRef->Extent().height), 0.1f, 10.0f);
    ubo.proj[1][1] *= -1;
    memcpy(BufferManager::uniformBuffers.UniformBuffersMapped()[currentImage], &ubo, sizeof(ubo));
}

void Lumine::LumineRenderer::recreateCommandBuffers()
{
}
